from pydantic import BaseModel, Field
from typing import List, Dict, Any
from time import sleep

import openai

MAX_RETRY = 5


class CompletionFailure(Exception):
    pass


class OpenaiCompatibleModel(BaseModel):
    api_base: str
    api_key: str
    api_model_name: str


class ChatModel(OpenaiCompatibleModel):
    temperature: float = 0.6
    system_prompt: str = ""

    def getInstance(self) -> 'ChatInstance':
        return ChatInstance(model=self)


class ChatInstance(BaseModel):
    model: ChatModel
    messages: List[Dict[str, Any]] = Field(default_factory=list)
    temperature: float = 0.0

    def setSystemPrompt(self, system_prompt: str) -> None:
        pass

    def ask(self, question: str) -> str:
        conn = openai.OpenAI(
            base_url=self.model.api_base,
            api_key=self.model.api_key
        )
        self.messages.append({'role': 'user', 'content': question})

        nretry = 0
        model_name = self.model.api_model_name
        while nretry <= MAX_RETRY:
            try:
                resp = conn.chat.completions.create(
                    messages=self.messages,  # type: ignore
                    model=model_name
                )
                answer = resp.choices[0].message.content or ""
                break
            except openai.RateLimitError:
                sleep(10)
                nretry += 1
            except openai.PermissionDeniedError:
                raise CompletionFailure(f"当前服务提供商不允许使用该模型：{self.model.api_model_name}")
        else:
            raise CompletionFailure(f"重试超过{MAX_RETRY}次，仍然无法完成请求")

        return answer
